#include "com_andon_lcz_h264rtmp_pusher_PusherNative.h"

JNIEXPORT void JNICALL
Java_com_andon_lcz_h264rtmp_pusher_PusherNative_setVideoOptions
        (JNIEnv *, jobject, jint _w, jint _h, jint _bitrate, jint _fps) {
    x264_param_t param;
    y_len = _w * _h;
    u_v_len = y_len / 4;
    x264_param_default_preset(&param, "ultrafast", "zerolatency");
    // 设置视频质量
    param.i_level_idc = 51;
    param.i_csp = X264_CSP_I420;
    param.i_width = _w;
    param.i_height = _h;
    param.i_threads = 1;
    param.i_timebase_num = param.i_fps_num = _fps;
    param.i_timebase_den = param.i_fps_den = 1;
    param.i_keyint_max = _fps * 2;
    // CQP 恒定质量 CRF 恒定码率 ABR 平均码率
    param.rc.i_rc_method = X264_RC_ABR;
    // 设置码率
    param.rc.i_bitrate = _bitrate / 1000;
    param.rc.i_vbv_max_bitrate = param.rc.i_bitrate * 1.2;
    // 设置缓冲区
    param.rc.i_vbv_buffer_size = param.rc.i_bitrate;
    // 设置 0 使用 pts 做音视频同步
    param.b_vfr_input = 0;
    // pps: 68 (图像参数集) sps: 67 (序列参数集 profile, level, 宽高，颜色空间等信息)
    param.b_repeat_headers = 1;
    // 设置视频质量，只有 I，P 帧，编解码速度快
    x264_param_apply_profile(&param, "baseline");

    // 打开编码器
    x264 = x264_encoder_open(&param);
    if (!x264) {
        LOG_E("打开编码器失败");
        return;
    }
    x264_pic_in = (x264_picture_t *) malloc(sizeof(x264_picture_t));
    x264_picture_alloc(x264_pic_in, X264_CSP_I420, _w, _h);
    x264_pic_out = (x264_picture_t *) malloc(sizeof(x264_picture_t));
    x264_picture_init(x264_pic_out);
}

JNIEXPORT void JNICALL
Java_com_andon_lcz_h264rtmp_pusher_PusherNative_pushVideo
        (JNIEnv *env, jobject, jbyteArray array) {

    // nv21 y、v/u... 转换为 yuv y、u...、v...
    jbyte *data = env->GetByteArrayElements(array, NULL);
    memcpy(x264_pic_in->img.plane[0], data, y_len);
    jbyte *u = (jbyte *) x264_pic_in->img.plane[1];
    jbyte *v = (jbyte *) x264_pic_in->img.plane[2];
    for (int i = 0; i < u_v_len; i++) {
        *(u + i) = *(data + y_len + i * 2 + 1);
        *(v + i) = *(data + y_len + i * 2);
    }
    // NALU 数组，一幅画面的若干切片
    x264_nal_t *nal = NULL;
    int n_nal = -1;
    x264_picture_init(x264_pic_out);
    if (x264_encoder_encode(x264, &nal, &n_nal, x264_pic_in, x264_pic_out) < 0) {
        LOG_E("编码失败");
        return;
    }
    x264_pic_in->i_pts++;

    // x264 中 1f 固定值
    uint8_t sps[100];
    memset(sps, 0, 100);
    int sps_len;
    uint8_t pps[100];
    memset(pps, 0, 100);
    int pps_len;
    for (int i = 0; i < n_nal; i++) {
        if (nal[i].i_type == NAL_SPS) {
            // 67 & 1f = NAL_SPS
            // 前面4个字节是分隔符 0000 0001
            sps_len = nal[i].i_payload - 4;
            memcpy(sps, nal[i].p_payload + 4, sps_len);
        } else if (nal[i].i_type == NAL_PPS) {
            // 68 & 1f = NAL_PPS
            pps_len = nal[i].i_payload - 4;
            memcpy(pps, nal[i].p_payload + 4, pps_len);
            // 发送 sps + pps 头信息
            add_264_sequence_header(pps, sps, pps_len, sps_len);
        } else {
            // 普通帧数据
            add_264_body(nal[i].p_payload, nal[i].i_payload);
        }
    }
}

void add_264_sequence_header(uint8_t *pps, uint8_t *sps, int pps_len, int sps_len) {
    int body_size = pps_len + sps_len + 16;
    RTMPPacket *packet = (RTMPPacket *) malloc(sizeof(RTMPPacket));
    RTMPPacket_Alloc(packet, body_size);
    int i = 0;
    unsigned char *body = (unsigned char *) packet->m_body;
    body[i++] = 0x17;
    body[i++] = 0x00;
    body[i++] = 0x00;
    body[i++] = 0x00;
    body[i++] = 0x00;
    // 版本号
    body[i++] = 0x01;
    // profile
    body[i++] = sps[1];
    // 兼容性
    body[i++] = sps[2];
    // 设置 profile level "baseline"
    body[i++] = sps[3];
    // 0xff 包长数据所使用的字节数
    body[i++] = 0xff;
    // sps 个数
    body[i++] = 0xe1;
    // sps 长度
    body[i++] = (sps_len >> 8) & 0xff;
    body[i++] = sps_len & 0xff;
    // sps 内容
    memcpy(&body[i], sps, sps_len);
    i += sps_len;
    body[i++] = 0x01;
    // pps 长度
    body[i++] = (pps_len >> 8) & 0xff;
    body[i++] = pps_len & 0xff;
    // pps 内容
    memcpy(&body[i], pps, pps_len);
    packet->m_packetType = RTMP_PACKET_TYPE_VIDEO;
    packet->m_nBodySize = body_size;
    packet->m_nTimeStamp = 0;
    packet->m_hasAbsTimestamp = 0;
    packet->m_nChannel = 0x04;
    packet->m_headerType = RTMP_PACKET_SIZE_MEDIUM;

    // 将 RTMP 包放入推送队列
    put(packet);
}

// 分隔符有两种 00 00 00 01 和 00 00 01
void add_264_body(uint8_t *buf, int buf_len) {
    // 去除分隔符，首地址偏移
    if (buf[2] == 0x00) {
        buf += 4;
        buf_len -= 4;
    } else if (buf[2] == 0x01) {
        buf += 3;
        buf_len -= 3;
    }
    int body_size = buf_len + 9;
    RTMPPacket *packet = (RTMPPacket *) malloc(sizeof(RTMPPacket));
    RTMPPacket_Alloc(packet, body_size);
    char *body = packet->m_body;
    // 判断是否关键帧单元 buf[0] == 65
    int type = buf[0] & 0x1f;
    body[0] = 0x27;
    if (type == NAL_SLICE_IDR) {
        body[0] = 0x17;
    }
    int i = 1;
    body[i++] = 0x01;
    body[i++] = 0x00;
    body[i++] = 0x00;
    body[i++] = 0x00;
    // 写入 body 的长度
    body[i++] = (buf_len >> 24) & 0xff;
    body[i++] = (buf_len >> 16) & 0xff;
    body[i++] = (buf_len >> 8) & 0xff;
    body[i++] = buf_len & 0xff;
    // 写入 body 的内容
    memcpy(&body[i], buf, buf_len);
    packet->m_packetType = RTMP_PACKET_TYPE_VIDEO;
    packet->m_nBodySize = body_size;
    packet->m_nTimeStamp = RTMP_GetTime() - start_time;
    packet->m_hasAbsTimestamp = 0;
    packet->m_nChannel = 0x04;
    packet->m_headerType = RTMP_PACKET_SIZE_LARGE;
    // 放到推流队列
    put(packet);
    LOG_I("放入一帧视频")
}

void put(RTMPPacket *packet) {
    pthread_mutex_lock(&p_mutex);
    if (publishing) {
        queue.push(packet);
    }
    pthread_cond_signal(&p_cond);
    pthread_mutex_unlock(&p_mutex);
}

RTMPPacket *get() {
    pthread_mutex_lock(&p_mutex);
    if (queue.empty()) {
        pthread_cond_wait(&p_cond, &p_mutex);
    }
    RTMPPacket *packet = queue.front();
    queue.pop();
    pthread_mutex_unlock(&p_mutex);
    return packet;
}

JNIEXPORT void JNICALL
Java_com_andon_lcz_h264rtmp_pusher_PusherNative_startPush(JNIEnv *env, jobject instance,
                                                          jstring url_) {
    const char *url = env->GetStringUTFChars(url_, 0);
    path = (char *) malloc(strlen(url) + 1);
    memcpy(path, url, strlen(url));
    env->ReleaseStringUTFChars(url_, url);
    pthread_mutex_init(&p_mutex, NULL);
    pthread_cond_init(&p_cond, NULL);
    start_time = RTMP_GetTime();
    pthread_t p_tid;
    pthread_create(&p_tid, NULL, push_thread, NULL);
}

void *push_thread(void *arg) {
    publishing = 1;
    RTMP *rtmp = RTMP_Alloc();
    RTMP_Init(rtmp);
    rtmp->Link.timeout = 5;
    RTMP_SetupURL(rtmp, path);
    RTMP_EnableWrite(rtmp);
    if (!RTMP_Connect(rtmp, NULL)) {
        LOG_E("连接失败");
        goto end;
    }
    // 连接流
    RTMP_ConnectStream(rtmp, 0);
    add_aac_sequence_header();
    while (publishing) {
        RTMPPacket *packet = get();
        // 推流
        packet->m_nInfoField2 = rtmp->m_stream_id;
        // 1 代表使用上传队列
        bool rs = RTMP_SendPacket(rtmp, packet, 1);
        LOG_I("推送 packet->%d", rs);
        RTMPPacket_Free(packet);
        free(packet);
    }
    end:
    publishing = 0;
    free(path);
    RTMP_Close(rtmp);
    RTMP_Free(rtmp);
    pthread_exit(NULL);
}

JNIEXPORT void JNICALL
Java_com_andon_lcz_h264rtmp_pusher_PusherNative_pushAudio(JNIEnv *env, jobject instance,
                                                          jbyteArray buf_, jint buf_size) {
    jbyte *buf = env->GetByteArrayElements(buf_, NULL);
    uint8_t *out_buf = (uint8_t *) malloc(sizeof(max_output_bytes));
    int out_buf_size = faacEncEncode(audioHandle, (int32_t *) buf, buf_size / 2, out_buf,
                                     max_output_bytes);
    if (out_buf_size > 0) {
        add_aac_body(out_buf, out_buf_size);
    }
    env->ReleaseByteArrayElements(buf_, buf, 0);
    free(out_buf);
}

void add_aac_sequence_header() {
    if (!audioHandle) {
        return;
    }
    unsigned char *buf;
    u_long len; /* buf长度,一般是 2 */
    faacEncGetDecoderSpecificInfo(audioHandle, &buf, &len);
    RTMPPacket *packet = (RTMPPacket *) malloc(sizeof(RTMPPacket));
    RTMPPacket_Alloc(packet, len + 2);
    RTMPPacket_Reset(packet);
    unsigned char *body = (unsigned char *) packet->m_body;
    /*AF 00 + AAC RAW data*/
    body[0] = 0xaf;
    body[1] = 0x00;
    memcpy(&body[2], buf, len); /*spec_buf 是 AAC sequence header 数据*/
    packet->m_packetType = RTMP_PACKET_TYPE_AUDIO;
    packet->m_nBodySize = len + 2;
    packet->m_nChannel = 0x04;
    packet->m_hasAbsTimestamp = 0;
    packet->m_nTimeStamp = 0;
    packet->m_hasAbsTimestamp = 0;
    packet->m_headerType = RTMP_PACKET_SIZE_MEDIUM;
    put(packet);
    free(buf);
    LOG_I("放入音频编码信息")
}

void add_aac_body(uint8_t *buf, int buf_len) {
    int body_size = buf_len + 2;
    RTMPPacket *packet = (RTMPPacket *) malloc(sizeof(RTMPPacket));
    RTMPPacket_Alloc(packet, body_size);
    RTMPPacket_Reset(packet);
    char *body = packet->m_body;
    body[0] = 0xaf;
    body[1] = 0x01;
    memcpy(&body[2], buf, buf_len);
    packet->m_packetType = RTMP_PACKET_TYPE_AUDIO;
    packet->m_nBodySize = body_size;
    packet->m_nTimeStamp = RTMP_GetTime() - start_time;
    packet->m_hasAbsTimestamp = 0;
    packet->m_nChannel = 0x04;
    packet->m_headerType = RTMP_PACKET_SIZE_MEDIUM;
    // 放到推流队列
    put(packet);
    LOG_I("放入一帧音频")
}

JNIEXPORT void JNICALL
Java_com_andon_lcz_h264rtmp_pusher_PusherNative_setAudioOptions(JNIEnv *env, jobject instance,
                                                                jint _sampleRate, jint _channels) {
    uint channels = _channels;
    u_long sampleRate = _sampleRate;
    audioHandle = faacEncOpen(sampleRate, channels, &input_samples, &max_output_bytes);
    if (!audioHandle) {
        // 失败
        LOG_I("打开音频编码器失败")
    }
    faacEncConfigurationPtr configurationPtr = faacEncGetCurrentConfiguration(audioHandle);
    configurationPtr->mpegVersion = MPEG4;
    configurationPtr->allowMidside = 1;
    configurationPtr->aacObjectType = LOW;
    // 是否包含 ADTS 头
    configurationPtr->outputFormat = 0;
    // 消除爆破声，短暂的噪声
    configurationPtr->useTns = 1;
    configurationPtr->useLfe = 0;
    configurationPtr->shortctl = SHORTCTL_NORMAL;
    configurationPtr->inputFormat = FAAC_INPUT_16BIT;
    configurationPtr->quantqual = 100;
    configurationPtr->bandWidth = 0; // 频宽

    // 使上面的配置生效
    if (!faacEncSetConfiguration(audioHandle, configurationPtr)) {
        LOG_E("失败");
    }
    LOG_I("声音设置成功");
}

JNIEXPORT jint JNICALL
Java_com_andon_lcz_h264rtmp_pusher_PusherNative_getInputSamples(JNIEnv *env, jobject instance) {
    return input_samples;
}

JNIEXPORT void JNICALL
Java_com_andon_lcz_h264rtmp_pusher_PusherNative_stopPush(JNIEnv *env, jobject instance) {
    if(publishing) {
        pthread_mutex_lock(&p_mutex);
        publishing = 0;
        pthread_cond_signal(&p_cond);
        pthread_mutex_unlock(&p_mutex);
    }
}
